"""
基于 LangChain 的股票分析智能体基类
"""
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import RunnablePassthrough
from langchain_core.messages import SystemMessage, HumanMessage
import json
from abc import ABC, abstractmethod
from colorama import Fore, Style
import logging
from typing import Dict, Any, Optional
from pydantic import BaseModel, Field

logger = logging.getLogger(__name__)

class AgentConfig(BaseModel):
    """智能体配置模型"""
    name: str = Field(description="智能体名称")
    model: str = Field(description="使用的模型名称")
    app_key: str = Field(description="API密钥")
    base_url: str = Field(description="API基础URL")
    prompt_template: str = Field(description="提示词模板")
    temperature: float = Field(default=0.5, description="温度参数")
    max_tokens: int = Field(default=2000, description="最大输出长度")

class StockAgent(ABC):
    """基于 LangChain 的股票分析智能体基类"""
    
    def __init__(self, agent_name: str):
        self.agent_name = agent_name
        self.config = self._load_config(agent_name)
        self.llm = self._create_llm()
        self.chain = self._create_chain()

    def _load_config(self, agent_name: str) -> AgentConfig:
        """加载智能体配置"""
        try:
            with open("config.json", "r", encoding="utf-8") as f:
                config_data = json.load(f)
                agent_config = config_data["agents"][agent_name]
                return AgentConfig(**agent_config)
        except (FileNotFoundError, json.JSONDecodeError, KeyError) as e:
            raise ValueError(f"Failed to load config for agent '{agent_name}': {e}")

    def _create_llm(self) -> ChatOpenAI:
        """创建 LangChain LLM 实例"""
        return ChatOpenAI(
            model=self.config.model,
            openai_api_key=self.config.app_key,
            openai_api_base=self.config.base_url,
            temperature=self.config.temperature,
            max_tokens=self.config.max_tokens,
            request_timeout=180  # 增加超时时间到3分钟
        )

    def _create_system_prompt(self) -> str:
        """创建系统提示词，子类可以重写"""
        return f"你是一个专业的{self.config.name}，请提供准确、有价值的分析。"

    def _create_chain(self):
        """创建 LangChain 处理链"""
        # 创建提示词模板
        system_prompt = SystemMessagePromptTemplate.from_template(
            self._create_system_prompt()
        )
        human_prompt = HumanMessagePromptTemplate.from_template("{input}")
        
        chat_prompt = ChatPromptTemplate.from_messages([
            system_prompt,
            human_prompt
        ])
        
        # 创建处理链
        chain = (
            {"input": RunnablePassthrough()}
            | chat_prompt
            | self.llm
            | StrOutputParser()
        )
        
        return chain

    def _execute_with_langchain(self, input_text: str, max_retries: int = 3) -> str:
        """使用 LangChain 执行推理"""
        print(Fore.MAGENTA + f"\n    - [LangChain] 启动 {self.config.name} 处理链...")
        print(Style.DIM + Fore.WHITE + f"      Input: {input_text[:200]}...")

        for attempt in range(max_retries):
            try:
                # 限制输入长度
                if len(input_text) > 8000:
                    input_text = input_text[:8000] + "\n\n[注: 输入内容已截断以避免请求过大]"
                
                # 使用 LangChain 链执行
                result = self.chain.invoke(input_text)
                
                print(Fore.GREEN + f"    - [LangChain] {self.config.name} 处理完成")
                print(Style.DIM + Fore.WHITE + f"      Output: {result[:200]}...")
                return result
                
            except Exception as e:
                print(Fore.YELLOW + f"    - [LangChain] 第 {attempt + 1} 次尝试失败: {str(e)}")
                
                if attempt < max_retries - 1:
                    import time
                    wait_time = (attempt + 1) * 2
                    print(Fore.YELLOW + f"    - [Retry] 等待 {wait_time} 秒后重试...")
                    time.sleep(wait_time)
                else:
                    error_message = f"LangChain execution error: {e}"
                    print(Fore.RED + f"    - [LangChain] 所有重试失败，{self.config.name} 处理出错")
                    print(Style.DIM + Fore.RED + f"      Error: {str(e)}")
                    return error_message

    @abstractmethod
    def execute(self, **kwargs) -> str:
        """每个智能体必须实现的主执行方法"""
        pass

    def get_agent_info(self) -> Dict[str, Any]:
        """获取智能体信息"""
        return {
            "name": self.config.name,
            "model": self.config.model,
            "base_url": self.config.base_url,
            "agent_type": self.__class__.__name__
        }

